#!/usr/bin/env node

// Initialize memory configuration early
import { MemoryManager } from './memory-init.js';
MemoryManager.initialize();

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
  McpError,
  ErrorCode,
} from '@modelcontextprotocol/sdk/types.js';

import { ClassFinder } from './class-finder.js';
import { JavaDecompiler } from './decompiler.js';
import { LocalRepository } from './local-repository.js';
import { configManager, config } from './config.js';

class JavaDecompilerMCPServer {
  private server: Server;
  private classFinder: ClassFinder;
  private decompiler: JavaDecompiler;
  private localRepository: LocalRepository;

  constructor() {
    console.error(`[MCP Server] 初始化Java反编译MCP服务器...`);
    
    // Print configuration
    configManager.printConfig();
    
    this.server = new Server({
      name: 'java-decompiler-mcp',
      version: '1.0.0',
    });

    this.classFinder = new ClassFinder();
    this.decompiler = new JavaDecompiler();
    this.localRepository = new LocalRepository(this.classFinder);
    
    // Use configuration from config file
    this.localRepository.setLocalRepositoryPath(config.maven.localRepository);
    console.error(`[MCP Server] 使用配置的Maven仓库路径: ${config.maven.localRepository}`);
    
    // Initialize repository index
    console.error(`[MCP Server] 初始化仓库索引...`);
    this.localRepository.initializeIndex().catch(error => {
      console.error(`[MCP Server] 索引初始化失败: ${error}`);
    });

    this.setupHandlers();
  }

  private setupHandlers() {
    this.server.setRequestHandler(ListToolsRequestSchema, async () => {
      return {
        tools: [
          {
            name: 'search_classes',
            description: 'Search for Java classes in local Maven repository (returns class paths only, no source code)',
            inputSchema: {
              type: 'object',
              properties: {
                keyword: {
                  type: 'string',
                  description: 'Keyword to search for in class names (e.g., "Audit", "Service")'
                },
                limit: {
                  type: 'number',
                  description: 'Maximum number of results to return (default: 50, max: 200)'
                },
                repositoryPath: {
                  type: 'string',
                  description: 'Optional: custom Maven repository path'
                }
              },
              required: ['keyword']
            }
          },
          {
            name: 'decompile_class',
            description: 'Decompile a specific Java class by its full class name',
            inputSchema: {
              type: 'object',
              properties: {
                fullClassName: {
                  type: 'string',
                  description: 'Full class name (e.g., "com.example.service.UserService")'
                },
                jarPath: {
                  type: 'string',
                  description: 'Optional: specific JAR file path if known'
                },
                repositoryPath: {
                  type: 'string',
                  description: 'Optional: custom Maven repository path'
                }
              },
              required: ['fullClassName']
            }
          }
        ]
      };
    });

    this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
      const { name, arguments: args } = request.params;

      try {
        switch (name) {
          case 'search_classes':
            return await this.handleSearchClasses(args as {
              keyword: string;
              limit?: number;
              repositoryPath?: string;
            });
          
          case 'decompile_class':
            return await this.handleDecompileClass(args as {
              fullClassName: string;
              jarPath?: string;
              repositoryPath?: string;
            });
          
          // Legacy support with warning
          case 'find_class_in_local_repository':
            return await this.handleLegacyFindClass(args as {
              className: string;
              repositoryPath?: string;
            });

          default:
            throw new McpError(
              ErrorCode.MethodNotFound,
              `Unknown tool: ${name}`
            );
        }
      } catch (error) {
        throw new McpError(
          ErrorCode.InternalError,
          `Error executing tool ${name}: ${error}`
        );
      }
    });
  }


  // New method: Search classes (only returns metadata, no source code)
  private async handleSearchClasses(args: { keyword: string; limit?: number; repositoryPath?: string }) {
    console.error(`[MCP Server] 收到搜索类请求: ${args.keyword}`);
    
    const limit = Math.min(args.limit || 50, 200); // Default 50, max 200
    
    try {
      if (args.repositoryPath) {
        console.error(`[MCP Server] 使用自定义仓库路径: ${args.repositoryPath}`);
        this.localRepository.setLocalRepositoryPath(args.repositoryPath);
        await this.localRepository.initializeIndex();
      }

      console.error(`[MCP Server] 搜索关键字: ${args.keyword}, 限制: ${limit}`);
      const indexEntries = await this.localRepository.findClassesByName(args.keyword);
      const limitedEntries = indexEntries.slice(0, limit);
      
      console.error(`[MCP Server] 找到 ${indexEntries.length} 个匹配，返回 ${limitedEntries.length} 个`);

      const searchResults = limitedEntries.map(entry => ({
        fullName: entry.fullName,
        className: entry.className,
        packageName: entry.packageName,
        artifact: {
          groupId: entry.artifact.groupId,
          artifactId: entry.artifact.artifactId,
          version: entry.artifact.version,
          jarPath: entry.artifact.localPath
        }
      }));

      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              keyword: args.keyword,
              totalFound: indexEntries.length,
              returned: searchResults.length,
              results: searchResults
            }, null, 2)
          }
        ]
      };
    } catch (error) {
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              success: false,
              error: `Search failed: ${error}`,
              keyword: args.keyword
            })
          }
        ]
      };
    }
  }

  // New method: Decompile a specific class
  private async handleDecompileClass(args: { fullClassName: string; jarPath?: string; repositoryPath?: string }) {
    console.error(`[MCP Server] 收到反编译请求: ${args.fullClassName}`);
    
    try {
      if (args.repositoryPath) {
        this.localRepository.setLocalRepositoryPath(args.repositoryPath);
        await this.localRepository.initializeIndex();
      }

      let classInfo;
      let classBytes;

      if (args.jarPath) {
        // If JAR path is provided, use it directly
        console.error(`[MCP Server] 使用指定的JAR路径: ${args.jarPath}`);
        classInfo = {
          className: args.fullClassName.split('.').pop() || args.fullClassName,
          fullName: args.fullClassName,
          packageName: args.fullClassName.includes('.') ? 
            args.fullClassName.substring(0, args.fullClassName.lastIndexOf('.')) : '',
          jarPath: args.jarPath
        };
        classBytes = await this.classFinder.getClassSource(classInfo);
      } else {
        // Search for the class by exact full name
        console.error(`[MCP Server] 搜索精确匹配的类: ${args.fullClassName}`);
        const indexEntries = await this.localRepository.findClassesByFullName(args.fullClassName);
        
        if (indexEntries.length === 0) {
          return {
            content: [
              {
                type: 'text',
                text: JSON.stringify({
                  success: false,
                  error: `Class not found: ${args.fullClassName}. Please use 'search_classes' tool first to find the exact class name.`,
                  fullClassName: args.fullClassName
                })
              }
            ]
          };
        }

        // Use the first match
        const indexEntry = indexEntries[0];
        classInfo = {
          className: indexEntry.className,
          fullName: indexEntry.fullName,
          packageName: indexEntry.packageName,
          jarPath: indexEntry.artifact.localPath
        };
        classBytes = await this.classFinder.getClassSource(classInfo);
      }

      if (!classBytes) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: "Failed to extract class bytecode from JAR",
                fullClassName: args.fullClassName
              })
            }
          ]
        };
      }

      const decompileResult = await this.decompiler.decompileClass(classBytes, args.fullClassName);
      
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              success: decompileResult.success,
              fullClassName: args.fullClassName,
              jarPath: classInfo.jarPath,
              source: decompileResult.success ? decompileResult.source : undefined,
              error: !decompileResult.success ? decompileResult.error : undefined
            }, null, 2)
          }
        ]
      };
    } catch (error) {
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              success: false,
              error: `Decompilation failed: ${error}`,
              fullClassName: args.fullClassName
            })
          }
        ]
      };
    }
  }

  // Legacy method with deprecation warning
  private async handleLegacyFindClass(args: { className: string; repositoryPath?: string }) {
    console.error(`[MCP Server] ⚠️ Legacy API called: find_class_in_local_repository`);
    
    // Check if it looks like a full class name
    if (!args.className.includes('.')) {
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              warning: "This tool is deprecated. Please use the new two-step process:",
              steps: [
                "1. Use 'search_classes' to find classes by keyword",
                "2. Use 'decompile_class' with the full class name to get source code"
              ],
              hint: `For example, search for '${args.className}' first, then decompile specific class like 'com.example.${args.className}'`
            }, null, 2)
          }
        ]
      };
    }

    // If it's a full class name, redirect to decompile_class
    return this.handleDecompileClass({
      fullClassName: args.className,
      repositoryPath: args.repositoryPath
    });
  }

  private async handleFindClassInLocalRepository(args: { className: string; repositoryPath?: string }) {
    console.error(`[MCP Server] 收到在本地仓库中查找类请求: ${args.className}`);
    
    try {
      // Set custom repository path if provided
      if (args.repositoryPath) {
        console.error(`[MCP Server] 使用自定义仓库路径: ${args.repositoryPath}`);
        this.localRepository.setLocalRepositoryPath(args.repositoryPath);
        // Re-initialize index for new repository path
        await this.localRepository.initializeIndex();
      }

      // Use index-based search for much better performance
      console.error(`[MCP Server] 使用索引快速搜索类: ${args.className}`);
      const indexEntries = await this.localRepository.findClassesByName(args.className);
      console.error(`[MCP Server] 索引搜索完成，找到 ${indexEntries.length} 个匹配的类`);

      const foundClasses: any[] = [];

      // Process each found class and decompile it
      for (const indexEntry of indexEntries) {
        try {
          console.error(`[MCP Server] 反编译类: ${indexEntry.fullName}`);
          
          // Create a temporary ClassInfo object for decompilation
          const classInfo = {
            className: indexEntry.className,
            fullName: indexEntry.fullName,
            packageName: indexEntry.packageName,
            jarPath: indexEntry.artifact.localPath
          };
          
          // Get class bytes and decompile
          const classBytes = await this.classFinder.getClassSource(classInfo);
          if (classBytes) {
            const decompileResult = await this.decompiler.decompileClass(classBytes, indexEntry.fullName);
            
            foundClasses.push({
              className: indexEntry.className,
              fullName: indexEntry.fullName,
              packageName: indexEntry.packageName,
              artifact: {
                groupId: indexEntry.artifact.groupId,
                artifactId: indexEntry.artifact.artifactId,
                version: indexEntry.artifact.version,
                localPath: indexEntry.artifact.localPath
              },
              decompiled: decompileResult.success,
              source: decompileResult.success ? decompileResult.source : decompileResult.error
            });
            
            console.error(`[MCP Server] 类 ${indexEntry.fullName} 反编译${decompileResult.success ? '成功' : '失败'}`);
          } else {
            console.error(`[MCP Server] 无法获取类字节码: ${indexEntry.fullName}`);
            foundClasses.push({
              className: indexEntry.className,
              fullName: indexEntry.fullName,
              packageName: indexEntry.packageName,
              artifact: {
                groupId: indexEntry.artifact.groupId,
                artifactId: indexEntry.artifact.artifactId,
                version: indexEntry.artifact.version,
                localPath: indexEntry.artifact.localPath
              },
              decompiled: false,
              source: "Failed to extract class bytecode from JAR"
            });
          }
        } catch (error) {
          console.error(`[MCP Server] 处理类失败 ${indexEntry.fullName}: ${error}`);
          foundClasses.push({
            className: indexEntry.className,
            fullName: indexEntry.fullName,
            packageName: indexEntry.packageName,
            artifact: {
              groupId: indexEntry.artifact.groupId,
              artifactId: indexEntry.artifact.artifactId,
              version: indexEntry.artifact.version,
              localPath: indexEntry.artifact.localPath
            },
            decompiled: false,
            source: `Decompilation failed: ${error}`
          });
        }
      }
      
      console.error(`[MCP Server] 索引搜索完成，成功处理 ${foundClasses.length} 个类`);

      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              className: args.className,
              repositoryPath: args.repositoryPath || this.localRepository.getLocalRepositoryPath(),
              totalArtifactsScanned: await this.localRepository.getIndexedArtifactCount(),
              totalClassesIndexed: await this.localRepository.getIndexedClassCount(),
              matchesFound: foundClasses.length,
              indexReady: this.localRepository.isIndexReady(),
              classes: foundClasses
            }, null, 2)
          }
        ]
      };

    } catch (error) {
      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              success: false,
              error: `Failed to find class in local repository: ${error}`,
              className: args.className,
              indexReady: this.localRepository.isIndexReady()
            })
          }
        ]
      };
    }
  }

  async run() {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
    console.error('Java Decompiler MCP server running on stdio');
  }
}

const server = new JavaDecompilerMCPServer();
server.run().catch(console.error);